import os
from openai import OpenAI
from dotenv import load_dotenv
from loguru import logger
from typing import List, Dict, Optional
from .knowledge_base import MedicalKnowledgeBase

load_dotenv()

class DeepSeekFinetuner:
    def __init__(self, use_knowledge_base: bool = True):
        self.api_key = os.getenv('DEEPSEEK_API_KEY')
        if not self.api_key or self.api_key == 'your_api_key_here':
            raise ValueError("DeepSeek API密钥未正确配置")
        
        self.client = OpenAI(
            base_url="https://api.deepseek.com/v1",
            api_key=self.api_key
        )
        
        # 初始化知识库
        self.use_knowledge_base = use_knowledge_base
        if use_knowledge_base:
            try:
                self.knowledge_base = MedicalKnowledgeBase()
                logger.info("成功初始化医疗知识库")
            except Exception as e:
                logger.error(f"初始化医疗知识库失败: {str(e)}")
                self.use_knowledge_base = False

    def api_call(self, prompt: str, department: str = None) -> dict:
        try:
            # 检索知识库获取相关信息
            retrieved_docs = []
            sources = []
            
            if self.use_knowledge_base:
                try:
                    logger.info(f"开始检索相关知识，科室：{department}")
                    # 加载相关科室的数据
                    self.knowledge_base.load_data(department)
                    
                    # 检索相关知识
                    retrieved_docs = self.knowledge_base.search(prompt, department)
                    sources = self.knowledge_base.get_sources(retrieved_docs)
                    
                    logger.info(f"知识库检索到{len(retrieved_docs)}条相关信息")
                except Exception as e:
                    logger.error(f"知识库检索失败: {str(e)}")
            
            # 构建增强提示
            enhanced_prompt = prompt
            
            if retrieved_docs:
                # 将检索到的知识添加到提示中
                logger.debug(f"检索到的知识: {retrieved_docs}")
                knowledge_context = "\n\n参考以下医学知识：\n"
                for i, doc in enumerate(retrieved_docs):
                    knowledge_context += f"[{i+1}] {doc['answer']}\n"
                
                enhanced_prompt = knowledge_context + "\n" + prompt
                logger.debug(f"增强后的提示: {enhanced_prompt[:200]}...")
                logger.debug(f"增强后的提示  all content : {enhanced_prompt} ")
            else:
                logger.info("未检索到相关知识，仅使用原始提示")
            # 调用API生成回答
            logger.info("开始调用DeepSeek API生成回答")
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[{"role": "user", "content": enhanced_prompt}],
                temperature=0.3,
                max_tokens=2048
            )
            logger.info("DeepSeek API调用成功")
            logger.info("DeepSeek response : " + str(response))

            return {
                "choices": [{
                    "message": {
                        "content": response.choices[0].message.content
                    }
                }],
                "sources": sources
            }
        except Exception as e:
            logger.error(f"API调用失败: {str(e)}")
            raise ValueError("模型服务调用异常，请检查网络连接和API配置")